from typing import List, AnyStr, Union

import numpy as np
from bert_serving.client import BertClient
from sklearn.preprocessing import normalize

from config import *

bert_client = BertClient(ip=BERT_HOST, port=BERT_PORT, check_length=False)


def sentence_embedding(text: Union[AnyStr, List[AnyStr]]) -> List[List[float]]:
    """
    Sentence Embedding
    :param text: 文本, 支持str类型(一个文本), list类型(多个文本)
    :return: 二维数组, shape: (n, BERT_LENGTH), BERT_LENGTH默认为768, n为文本数目
    """
    input_text = None
    if isinstance(text, str):
        input_text = [text]
    elif isinstance(text, list):
        input_text = text
    else:
        raise TypeError('只支持List[Str]或Str')
    raw_embedding = bert_client.encode(input_text)
    embedding = normalize(raw_embedding)
    print(type(embedding), embedding.shape)
    return embedding.tolist()

